package com.tianji.aigc.service.impl;

import cn.hutool.core.date.DateUtil;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.vo.ChatEventVO;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

/**
 * @Description: TODO
 * @Author: lemon
 * @CreateTime: 2025-06-07  16:40
 * @Version: 1.0
 */
@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final ChatClient chatClient;

    private final SystemPromptConfig systemPromptConfig;

    private final StringRedisTemplate redisTemplate;

    private final ChatMemory chatMemory;

    //生成状态前缀
    private static final String GENERATE_STATUS_PREFIX = "generate_status:";

    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        //获取对话id
        String conversationId = ChatService.getConversationId(sessionId);
        //大模型输出内容的缓存器，用于在输出中断后存储数据
        StringBuilder outputBuilder = new StringBuilder();
        return chatClient.prompt()
                .user(question)
                .system(promptSystem -> promptSystem
                        .text(systemPromptConfig.getChatSystemMessage().get())
                        .param("now", DateUtil.now()))
                .advisors(advisor -> advisor.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId))
                .stream()
                .chatResponse()
                .doFirst(() -> redisTemplate.opsForValue().set(getGenerateKey(sessionId), "true"))   //第一次输出时标记为生成状态 true
                .doOnCancel(() -> saveStopHistoryRecord(conversationId, outputBuilder.toString()))   //终止对话时，保存中断后的大模型输出数据
                .doOnComplete(() -> redisTemplate.delete(getGenerateKey(sessionId)))  //输出结束时清除标记
                // 输出过程中，判断是否正在输出，如果正在输出，则继续输出，否则结束输出
                .takeWhile(s -> redisTemplate.opsForValue().get(getGenerateKey(sessionId)) != null)
                .map(chatResponse -> {
                    //获取大模型输出的文本内容
                    String text = chatResponse.getResult().getOutput().getText();
                    //追加到输出内容中
                    outputBuilder.append(text);
                    //将输出内容封装为ChatEventVO对象 -- 统一格式
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.just(ChatEventVO.builder()
                        .eventType(ChatEventTypeEnum.STOP.getValue())
                        .build()));  //停止事件
    }


    /**
     * 保存停止输出的记录
     *
     * @param conversationId 会话id
     * @param content        大模型输出的内容
     */
    private void saveStopHistoryRecord(String conversationId, String content) {
        chatMemory.add(conversationId, new AssistantMessage(content));
    }

    /**
     * 停止生成
     *
     * @param sessionId
     */
    @Override
    public void stop(String sessionId) {
        redisTemplate.delete(getGenerateKey(sessionId));
    }

    public String getGenerateKey(String sessionId) {
        return GENERATE_STATUS_PREFIX + sessionId;
    }
}
